"""
102受试者SSVEP数据集主分类器
整合数据加载、预处理和FBCCA分类
"""

import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from dataset_loader import SSVEP102DatasetLoader
from fbcca_classifier import FBCCAClassifier, SSVEPPreprocessor, evaluate_classifier
from typing import List, Tuple
import os
from datetime import datetime

# 设置matplotlib中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False    # 用来正常显示负号

def create_results_folder() -> str:
    """创建结果保存文件夹"""
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    folder_name = f"ssvep_classification_results_{timestamp}"
    
    if not os.path.exists(folder_name):
        os.makedirs(folder_name)
        print(f"创建结果文件夹: {folder_name}")
    
    return folder_name

def determine_ssvep_frequencies() -> List[float]:
    """
    确定SSVEP频率列表
    从stimulation_information.pdf中获取的实际频率值
    """
    # 从stimulation_information.pdf中读取的实际频率
    # 按字符顺序: 1,2,3,4,5,6,7,8,9,0,*,# 对应目标索引 0-11
    frequencies = [9.25, 11.25, 13.25, 9.75, 11.75, 13.75, 10.25, 12.25, 14.25, 10.75, 12.75, 14.75]
    character_labels = ['1', '2', '3', '4', '5', '6', '7', '8', '9', '0', '*', '#']
    
    print(f"使用SSVEP频率: {frequencies}")
    print(f"对应字符标签: {character_labels}")
    print("频率来源: stimulation_information.pdf")
    
    return frequencies

def plot_confusion_matrix(cm: np.ndarray, class_names: List[str], title: str = "混淆矩阵", 
                          save_path: str = None):
    """绘制混淆矩阵"""
    plt.figure(figsize=(12, 10))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
                xticklabels=class_names, yticklabels=class_names)
    plt.title(title, fontsize=16)
    plt.xlabel('预测标签', fontsize=14)
    plt.ylabel('真实标签', fontsize=14)
    plt.xticks(rotation=45)
    plt.yticks(rotation=0)
    plt.tight_layout()
    
    if save_path:
        plt.savefig(os.path.join(save_path, f"{title.replace(' ', '_')}_confusion_matrix.png"), 
                   dpi=300, bbox_inches='tight')
        print(f"混淆矩阵已保存: {os.path.join(save_path, f'{title.replace(' ', '_')}_confusion_matrix.png')}")
    
    plt.show()

def plot_accuracy_by_subject(subject_accuracies: dict, title: str = "各受试者分类准确率", 
                            save_path: str = None):
    """绘制各受试者准确率"""
    subjects = list(subject_accuracies.keys())
    accuracies = list(subject_accuracies.values())
    
    plt.figure(figsize=(15, 8))
    bars = plt.bar(range(len(subjects)), accuracies, alpha=0.7, color='skyblue', edgecolor='navy')
    plt.xlabel('受试者ID', fontsize=14)
    plt.ylabel('准确率', fontsize=14)
    plt.title(title, fontsize=16)
    plt.xticks(range(len(subjects)), [f'S{s:03d}' for s in subjects], rotation=45)
    
    # 添加平均线
    mean_acc = np.mean(accuracies)
    plt.axhline(y=mean_acc, color='r', linestyle='--', linewidth=2,
                label=f'平均准确率: {mean_acc:.3f}')
    
    # 在柱状图上显示数值
    for i, (bar, acc) in enumerate(zip(bars, accuracies)):
        plt.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.01, 
                f'{acc:.3f}', ha='center', va='bottom', fontsize=8)
    
    plt.legend(fontsize=12)
    plt.grid(axis='y', alpha=0.3)
    plt.ylim(0, 1.1)
    plt.tight_layout()
    
    if save_path:
        plt.savefig(os.path.join(save_path, f"{title.replace(' ', '_')}_accuracy_by_subject.png"), 
                   dpi=300, bbox_inches='tight')
        print(f"受试者准确率图已保存: {os.path.join(save_path, f'{title.replace(' ', '_')}_accuracy_by_subject.png')}")
    
    plt.show()

def plot_performance_summary(results: dict, title: str = "性能总结", save_path: str = None):
    """绘制性能总结图（准确率和ITR）"""
    if 'subject_results' in results:
        # 受试者依赖评估结果
        subjects = list(results['subject_results'].keys())
        accuracies = [results['subject_results'][sid]['accuracy'] for sid in subjects]
        itrs = [results['subject_results'][sid]['itr'] for sid in subjects]
        
        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(15, 12))
        
        # 准确率子图
        bars1 = ax1.bar(range(len(subjects)), accuracies, alpha=0.7, color='lightcoral', edgecolor='darkred')
        ax1.set_xlabel('受试者ID', fontsize=12)
        ax1.set_ylabel('准确率', fontsize=12)
        ax1.set_title('各受试者分类准确率', fontsize=14)
        ax1.set_xticks(range(len(subjects)))
        ax1.set_xticklabels([f'S{s:03d}' for s in subjects], rotation=45)
        ax1.axhline(y=np.mean(accuracies), color='red', linestyle='--', 
                   label=f'平均: {np.mean(accuracies):.3f}')
        ax1.legend()
        ax1.grid(axis='y', alpha=0.3)
        ax1.set_ylim(0, 1.1)
        
        # ITR子图
        bars2 = ax2.bar(range(len(subjects)), itrs, alpha=0.7, color='lightgreen', edgecolor='darkgreen')
        ax2.set_xlabel('受试者ID', fontsize=12)
        ax2.set_ylabel('ITR (bits/min)', fontsize=12)
        ax2.set_title('各受试者信息传输率', fontsize=14)
        ax2.set_xticks(range(len(subjects)))
        ax2.set_xticklabels([f'S{s:03d}' for s in subjects], rotation=45)
        ax2.axhline(y=np.mean(itrs), color='green', linestyle='--', 
                   label=f'平均: {np.mean(itrs):.2f}')
        ax2.legend()
        ax2.grid(axis='y', alpha=0.3)
        
        plt.suptitle(title, fontsize=16)
        plt.tight_layout()
        
    else:
        # 受试者无关评估结果
        fig, ax = plt.subplots(1, 1, figsize=(8, 6))
        metrics = ['准确率', 'ITR (bits/min)']
        values = [results['accuracy'], results['itr']]
        colors = ['skyblue', 'lightgreen']
        
        bars = ax.bar(metrics, values, color=colors, alpha=0.7, edgecolor='navy')
        ax.set_ylabel('数值')
        ax.set_title(title, fontsize=16)
        
        # 在柱状图上显示数值
        for i, (bar, value) in enumerate(zip(bars, values)):
            if i == 1:  # ITR 值
                ax.text(bar.get_x() + bar.get_width()/2, bar.get_height() + bar.get_height()*0.01, 
                       f'{value:.2f}', ha='center', va='bottom', fontsize=12)
            else:  # 准确率
                ax.text(bar.get_x() + bar.get_width()/2, bar.get_height() + bar.get_height()*0.01, 
                       f'{value:.3f}', ha='center', va='bottom', fontsize=12)
        
        plt.tight_layout()
    
    if save_path:
        plt.savefig(os.path.join(save_path, f"{title.replace(' ', '_')}_performance_summary.png"), 
                   dpi=300, bbox_inches='tight')
        print(f"性能总结图已保存: {os.path.join(save_path, f'{title.replace(' ', '_')}_performance_summary.png')}")
    
    plt.show()

def plot_accuracy_distribution(results: dict, title: str = "准确率分布", save_path: str = None):
    """绘制准确率分布直方图"""
    if 'subject_results' not in results:
        return
    
    accuracies = [results['subject_results'][sid]['accuracy'] for sid in results['subject_results']]
    
    plt.figure(figsize=(10, 6))
    plt.hist(accuracies, bins=15, alpha=0.7, color='lightblue', edgecolor='navy')
    plt.axvline(np.mean(accuracies), color='red', linestyle='--', linewidth=2,
                label=f'平均值: {np.mean(accuracies):.3f}')
    plt.axvline(np.median(accuracies), color='green', linestyle='--', linewidth=2,
                label=f'中位数: {np.median(accuracies):.3f}')
    
    plt.xlabel('准确率', fontsize=14)
    plt.ylabel('受试者数量', fontsize=14)
    plt.title(title, fontsize=16)
    plt.legend(fontsize=12)
    plt.grid(axis='y', alpha=0.3)
    
    # 添加统计信息
    stats_text = f'样本数: {len(accuracies)}\n标准差: {np.std(accuracies):.3f}\n最小值: {min(accuracies):.3f}\n最大值: {max(accuracies):.3f}'
    plt.text(0.02, 0.98, stats_text, transform=plt.gca().transAxes, 
             verticalalignment='top', bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8))
    
    plt.tight_layout()
    
    if save_path:
        plt.savefig(os.path.join(save_path, f"{title.replace(' ', '_')}_distribution.png"), 
                   dpi=300, bbox_inches='tight')
        print(f"准确率分布图已保存: {os.path.join(save_path, f'{title.replace(' ', '_')}_distribution.png')}")
    
    plt.show()

def save_results_summary(results: dict, save_path: str, evaluation_type: str):
    """保存结果摘要到文本文件"""
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    
    with open(os.path.join(save_path, "results_summary.txt"), 'w', encoding='utf-8') as f:
        f.write(f"SSVEP分类结果摘要\n")
        f.write(f"生成时间: {timestamp}\n")
        f.write(f"评估类型: {evaluation_type}\n")
        f.write("="*50 + "\n\n")
        
        if evaluation_type == "受试者依赖评估":
            f.write(f"测试受试者数: {results['n_subjects']}\n")
            f.write(f"平均准确率: {results['mean_accuracy']:.4f} ± {results['std_accuracy']:.4f}\n")
            f.write(f"平均ITR: {results['mean_itr']:.2f} ± {results['std_itr']:.2f} bits/min\n")
            f.write(f"最高准确率: {max([results['subject_results'][sid]['accuracy'] for sid in results['subject_results']]):.4f}\n")
            f.write(f"最低准确率: {min([results['subject_results'][sid]['accuracy'] for sid in results['subject_results']]):.4f}\n\n")
            
            f.write("各受试者详细结果:\n")
            f.write("-"*30 + "\n")
            for sid in sorted(results['subject_results'].keys()):
                result = results['subject_results'][sid]
                f.write(f"受试者 S{sid:03d}: 准确率={result['accuracy']:.4f}, ITR={result['itr']:.2f} bits/min\n")
                
        else:  # 受试者无关评估
            f.write(f"训练受试者数: {results['n_train_subjects']}\n")
            f.write(f"测试受试者数: {results['n_test_subjects']}\n")
            f.write(f"准确率: {results['accuracy']:.4f}\n")
            f.write(f"ITR: {results['itr']:.2f} bits/min\n")
            f.write(f"训练受试者: {results['train_subjects']}\n")
            f.write(f"测试受试者: {results['test_subjects']}\n")
    
    print(f"结果摘要已保存: {os.path.join(save_path, 'results_summary.txt')}")

def subject_dependent_evaluation(loader: SSVEP102DatasetLoader, 
                                classifier_params: dict,
                                subject_ids: List[int] = None,
                                electrode_type: int = 1) -> dict:
    """
    受试者依赖评估 (每个受试者单独训练测试)
    
    参数:
    loader: 数据加载器
    classifier_params: 分类器参数
    subject_ids: 受试者ID列表，None表示使用所有受试者
    electrode_type: 电极类型 (1: 湿电极, 2: 干电极)
    
    返回:
    results: 评估结果字典
    """
    if subject_ids is None:
        subject_ids = list(range(1, 21))  # 默认测试前20个受试者
    
    print(f"开始受试者依赖评估...")
    print(f"测试受试者: {subject_ids}")
    print(f"电极类型: {'湿电极' if electrode_type == 1 else '干电极'}")
    
    # 创建预处理器
    preprocessor = SSVEPPreprocessor(fs=250)
    
    subject_results = {}
    all_accuracies = []
    all_itrs = []
    
    for subject_id in subject_ids:
        print(f"\n{'='*50}")
        print(f"处理受试者 {subject_id}...")
        
        # 加载受试者数据
        subject_data = loader.load_subject_data(subject_id)
        if subject_data is None:
            print(f"跳过受试者 {subject_id} (数据加载失败)")
            continue
        
        # 提取试验数据 (使用刺激窗口 0.5-2.5s)
        trials, labels = loader.extract_trials(
            subject_data, 
            electrode_type=electrode_type,
            time_window=(0.5, 2.5)
        )
        
        # 数据预处理
        trials = preprocessor.preprocess(trials, apply_bandpass=True, apply_baseline=True)
        
        # 分割训练测试数据 (前8个区块训练，后2个区块测试)
        X_train, X_test, y_train, y_test = loader.split_train_test(
            trials, labels, 
            train_blocks=list(range(8)), 
            test_blocks=list(range(8, 10))
        )
        
        # 创建并训练分类器
        classifier = FBCCAClassifier(**classifier_params)
        classifier.fit(X_train, y_train)
        
        # 评估分类器
        character_labels = ['1', '2', '3', '4', '5', '6', '7', '8', '9', '0', '*', '#']
        class_names = [f"{char}({freq}Hz)" for char, freq in zip(character_labels, classifier.ssvep_freqs)]
        results = evaluate_classifier(classifier, X_test, y_test, class_names)
        
        # 保存结果
        subject_results[subject_id] = results
        all_accuracies.append(results['accuracy'])
        all_itrs.append(results['itr'])
        
        print(f"受试者 {subject_id} 结果:")
        print(f"  准确率: {results['accuracy']:.4f}")
        print(f"  ITR: {results['itr']:.2f} bits/min")
    
    # 计算总体统计
    overall_results = {
        'subject_results': subject_results,
        'mean_accuracy': np.mean(all_accuracies),
        'std_accuracy': np.std(all_accuracies),
        'mean_itr': np.mean(all_itrs),
        'std_itr': np.std(all_itrs),
        'n_subjects': len(all_accuracies)
    }
    
    print(f"\n{'='*50}")
    print(f"受试者依赖评估总体结果:")
    print(f"  测试受试者数: {overall_results['n_subjects']}")
    print(f"  平均准确率: {overall_results['mean_accuracy']:.4f} ± {overall_results['std_accuracy']:.4f}")
    print(f"  平均ITR: {overall_results['mean_itr']:.2f} ± {overall_results['std_itr']:.2f} bits/min")
    print(f"  最高准确率: {max(all_accuracies):.4f}")
    print(f"  最低准确率: {min(all_accuracies):.4f}")
    
    return overall_results

def subject_independent_evaluation(loader: SSVEP102DatasetLoader,
                                  classifier_params: dict,
                                  train_subjects: List[int],
                                  test_subjects: List[int],
                                  electrode_type: int = 1) -> dict:
    """
    受试者无关评估 (跨受试者训练测试)
    
    参数:
    loader: 数据加载器
    classifier_params: 分类器参数
    train_subjects: 训练受试者ID列表
    test_subjects: 测试受试者ID列表
    electrode_type: 电极类型
    
    返回:
    results: 评估结果字典
    """
    print(f"开始受试者无关评估...")
    print(f"训练受试者: {train_subjects}")
    print(f"测试受试者: {test_subjects}")
    print(f"电极类型: {'湿电极' if electrode_type == 1 else '干电极'}")
    
    # 创建预处理器
    preprocessor = SSVEPPreprocessor(fs=250)
    
    # 批量加载训练数据
    print("加载训练数据...")
    X_train_all, y_train_all, valid_train_subjects = loader.batch_load_subjects(
        train_subjects, electrode_type, time_window=(0.5, 2.5)
    )
    
    if len(X_train_all) == 0:
        print("无法加载训练数据")
        return {}
    
    # 预处理训练数据
    X_train_all = preprocessor.preprocess(X_train_all, apply_bandpass=True, apply_baseline=True)
    
    # 批量加载测试数据
    print("加载测试数据...")
    X_test_all, y_test_all, valid_test_subjects = loader.batch_load_subjects(
        test_subjects, electrode_type, time_window=(0.5, 2.5)
    )
    
    if len(X_test_all) == 0:
        print("无法加载测试数据")
        return {}
    
    # 预处理测试数据
    X_test_all = preprocessor.preprocess(X_test_all, apply_bandpass=True, apply_baseline=True)
    
    # 创建并训练分类器
    classifier = FBCCAClassifier(**classifier_params)
    classifier.fit(X_train_all, y_train_all)
    
    # 评估分类器
    character_labels = ['1', '2', '3', '4', '5', '6', '7', '8', '9', '0', '*', '#']
    class_names = [f"{char}({freq}Hz)" for char, freq in zip(character_labels, classifier.ssvep_freqs)]
    results = evaluate_classifier(classifier, X_test_all, y_test_all, class_names)
    
    # 添加受试者信息
    results['train_subjects'] = valid_train_subjects
    results['test_subjects'] = valid_test_subjects
    results['n_train_subjects'] = len(valid_train_subjects)
    results['n_test_subjects'] = len(valid_test_subjects)
    
    print(f"\n受试者无关评估结果:")
    print(f"  训练受试者数: {results['n_train_subjects']}")
    print(f"  测试受试者数: {results['n_test_subjects']}")
    print(f"  准确率: {results['accuracy']:.4f}")
    print(f"  ITR: {results['itr']:.2f} bits/min")
    
    return results

def main():
    """主函数"""
    print("102受试者SSVEP数据集FBCCA分类器")
    print("="*50)
    
    # 创建结果保存文件夹
    results_folder = create_results_folder()
    
    # 创建数据加载器
    loader = SSVEP102DatasetLoader()
    
    # 确定SSVEP频率
    ssvep_freqs = determine_ssvep_frequencies()
    
    # 分类器参数
    classifier_params = {
        'fs': 250,
        'num_harmonics': 3,
        'num_fb': 5,
        'window_length': 2.0,
        'ssvep_freqs': ssvep_freqs
    }
    
    # 选择评估模式
    evaluation_mode = input("选择评估模式:\n1. 受试者依赖评估\n2. 受试者无关评估\n请输入数字 (1-2): ")
    
    if evaluation_mode == "1":
        # 受试者依赖评估
        print("\n选择要测试的受试者数量:")
        print("1. 前10个受试者 (快速测试)")
        print("2. 前20个受试者 (中等规模)")
        print("3. 前50个受试者 (大规模测试)")
        print("4. 自定义范围")
        
        choice = input("请输入选择 (1-4): ")
        
        if choice == "1":
            subject_ids = list(range(1, 11))  # 前10个受试者
        elif choice == "2":
            subject_ids = list(range(1, 21))  # 前20个受试者
        elif choice == "3":
            subject_ids = list(range(1, 51))  # 前50个受试者
        elif choice == "4":
            start_id = int(input("起始受试者ID (1-102): "))
            end_id = int(input("结束受试者ID (1-102): "))
            subject_ids = list(range(start_id, end_id + 1))
        else:
            print("无效选择，使用默认前10个受试者")
            subject_ids = list(range(1, 11))
        
        print(f"\n将测试受试者: S{subject_ids[0]:03d} 到 S{subject_ids[-1]:03d} (共{len(subject_ids)}个)")
        print(f"预估时间: 约{len(subject_ids) * 2}分钟")
        confirm = input("确认继续? (y/n): ")
        if confirm.lower() != 'y':
            print("取消测试")
            return
            
        results = subject_dependent_evaluation(loader, classifier_params, subject_ids)
        
        # 保存结果摘要
        save_results_summary(results, results_folder, "受试者依赖评估")
        
        # 绘制并保存结果图像
        if results['subject_results']:
            subject_accuracies = {sid: results['subject_results'][sid]['accuracy'] 
                                for sid in results['subject_results']}
            
            # 绘制各受试者准确率
            plot_accuracy_by_subject(subject_accuracies, save_path=results_folder)
            
            # 绘制性能总结图
            plot_performance_summary(results, title="受试者依赖评估性能总结", save_path=results_folder)
            
            # 绘制准确率分布图
            plot_accuracy_distribution(results, save_path=results_folder)
    
    elif evaluation_mode == "2":
        # 受试者无关评估
        print("\n选择受试者无关评估规模:")
        print("1. 小规模 (3个训练, 2个测试)")
        print("2. 中规模 (10个训练, 5个测试)")
        print("3. 大规模 (20个训练, 10个测试)")
        print("4. 自定义")
        
        choice = input("请输入选择 (1-4): ")
        
        if choice == "1":
            train_subjects = list(range(1, 4))   # 前3个受试者训练
            test_subjects = list(range(4, 6))    # 后2个受试者测试
        elif choice == "2":
            train_subjects = list(range(1, 11))  # 前10个受试者训练
            test_subjects = list(range(11, 16))  # 后5个受试者测试
        elif choice == "3":
            train_subjects = list(range(1, 21))  # 前20个受试者训练
            test_subjects = list(range(21, 31))  # 后10个受试者测试
        elif choice == "4":
            train_start = int(input("训练集起始ID (1-102): "))
            train_end = int(input("训练集结束ID (1-102): "))
            test_start = int(input("测试集起始ID (1-102): "))
            test_end = int(input("测试集结束ID (1-102): "))
            train_subjects = list(range(train_start, train_end + 1))
            test_subjects = list(range(test_start, test_end + 1))
        else:
            print("无效选择，使用默认小规模")
            train_subjects = list(range(1, 4))
            test_subjects = list(range(4, 6))
        
        print(f"\n训练集: S{train_subjects[0]:03d}-S{train_subjects[-1]:03d} ({len(train_subjects)}个)")
        print(f"测试集: S{test_subjects[0]:03d}-S{test_subjects[-1]:03d} ({len(test_subjects)}个)")
        print(f"预估时间: 约{(len(train_subjects) + len(test_subjects)) * 1.5}分钟")
        confirm = input("确认继续? (y/n): ")
        if confirm.lower() != 'y':
            print("取消测试")
            return
            
        results = subject_independent_evaluation(
            loader, classifier_params, train_subjects, test_subjects
        )
        
        # 保存结果摘要
        save_results_summary(results, results_folder, "受试者无关评估")
        
        # 绘制并保存结果图像
        if 'confusion_matrix' in results:
            character_labels = ['1', '2', '3', '4', '5', '6', '7', '8', '9', '0', '*', '#']
            class_names = [f"{char}({freq}Hz)" for char, freq in zip(character_labels, ssvep_freqs)]
            
            # 绘制混淆矩阵
            plot_confusion_matrix(results['confusion_matrix'], class_names, 
                                title="受试者无关评估混淆矩阵", save_path=results_folder)
            
            # 绘制性能总结图
            plot_performance_summary(results, title="受试者无关评估性能总结", save_path=results_folder)
    
    else:
        print("无效的选择，退出程序")
        return
    
    print(f"\n分析完成!")
    print(f"所有结果已保存到文件夹: {results_folder}")
    print("包含以下文件:")
    print("- results_summary.txt: 结果摘要")
    print("- *.png: 各种可视化图像")

if __name__ == "__main__":
    main()
